#include "raytracer.h"
#include "math/MathLib.h"
#include "scene/scene.h"
#include "scene/objects.h"
#include "scene/texture.h"
#include "utility/image.h"

using rtMath::vec3;
using rtMath::vec4;
using rtMath::Ray;


struct TraycerSpeedup
{
	const RayTracer*	tracer;
	const Scene*		scene;
	Image*				screen;
	int					startLine;
	int					numThreads;
};

DWORD WINAPI ThreadWork(void* ptr)
{
	TraycerSpeedup* spdup = reinterpret_cast<TraycerSpeedup*>(ptr);

	const uint32 width = spdup->screen->GetWidth();
	const uint32 height = spdup->screen->GetHeight();
	const RayTracer* tracer = spdup->tracer;
	const Scene& scene = *(spdup->scene);
	Image* screen = spdup->screen;
	SceneCamera* camera = spdup->scene->GetCamera();

	const int startY = spdup->startLine;
	const int numThreads = spdup->numThreads;

	uint32 x, y;
	for (y = startY; y < height; y += numThreads)
		for (x = 0; x < width; ++x)
		{
			Ray centerRay = camera->ComputeRay(static_cast<float>(x), static_cast<float>(y),
											   static_cast<float>(width), static_cast<float>(height));
			vec4 resultColor = tracer->TraceRay(centerRay, scene, 0);
			screen->SetPixel(x, y, resultColor);
		}

	return 0;
}


RayTracer::RayTracer()
	: m_MaximumTraceDepth(s_MaxTraceDepth)
{
}
RayTracer::~RayTracer()
{
}


void RayTracer::SetMaximumTraceDepth(int maxTraceDepth)
{
	m_MaximumTraceDepth = s_MaxTraceDepth;
}

void RayTracer::Render(const Scene& scene, Image& screen)
{
	const uint32 width = screen.GetWidth();
	const uint32 height = screen.GetHeight();

	SceneCamera* camera = scene.GetCamera();
	camera->Setup(static_cast<float>(width), static_cast<float>(height));

	for (uint32 y = 0; y < height; ++y)
		for (uint32 x = 0; x < width; ++x)
		{
			Ray centerRay = camera->ComputeRay(static_cast<float>(x), static_cast<float>(y),
											   static_cast<float>(width), static_cast<float>(height));
			vec4 resultColor = this->TraceRay(centerRay, scene, 0);
			screen.SetPixel(x, y, resultColor);
		}
}

void RayTracer::RenderMP(const Scene& scene, Image& screen, int numThreads)
{
	TraycerSpeedup* speedups = new TraycerSpeedup[numThreads];
	HANDLE* nativeThreads = new HANDLE[numThreads];

	const uint32 width = screen.GetWidth();
	const uint32 height = screen.GetHeight();
	SceneCamera* camera = scene.GetCamera();
	camera->Setup(static_cast<float>(width), static_cast<float>(height));

	int i;
	for (i = 0; i < numThreads; ++i)
	{
		speedups[i].tracer = this;
		speedups[i].scene = &scene;
		speedups[i].screen = &screen;
		speedups[i].startLine = i;
		speedups[i].numThreads = numThreads;

		nativeThreads[i] = ::CreateThread(NULL, 0, static_cast<LPTHREAD_START_ROUTINE>(&ThreadWork), &speedups[i], 0, NULL);
	}

	::WaitForMultipleObjects(numThreads, nativeThreads, TRUE, INFINITE);

	for (i = 0; i < numThreads; ++i)
		::CloseHandle(nativeThreads[i]);

	SafeDeleteArray(nativeThreads);
	SafeDeleteArray(speedups);
}


rtMath::vec4 RayTracer::TraceRay(const rtMath::Ray& ray, const Scene& scene, int currentDepth) const
{
	Material defaultMaterial;
	defaultMaterial.SetDiffuse(vec3(1.0f, 1.0f, 1.0f));
	defaultMaterial.SetReflectivity(0.0f);

	const int numShapes = scene.GetShapesCount();
	const int numLights = scene.GetLightsCount();

	vec4 color(0.0f, 0.0f, 0.0f, 1);

	RayHitInfo closestHit;
	closestHit.distToHit = rtMath::MAX_FLOAT;
	int hitShapeId = -1;
	const Material* hitMaterial = NULL;

	for (int i = 0; i < numShapes; ++i)
	{
		RayHitInfo hitInfo;
		const SceneShape* shape = scene.GetShape(i);
		if (shape->RayCollisionCheck(ray, &hitInfo, closestHit.distToHit))
		{
			closestHit = hitInfo;
			hitShapeId = i;
		}
	}

	if (-1 == hitShapeId)
		return vec4(50.0f / 255.0f, 153.0f / 255.0f, 204.0f / 255.0f, 1);;

	const SceneShape* hitShape = scene.GetShape(hitShapeId);
	if (-1 != hitShape->GetMaterialID())
		hitMaterial = scene.GetMaterial(hitShape->GetMaterialID());

	if (!hitMaterial)
		hitMaterial = &defaultMaterial;

	vec3 materialDiffuse = hitMaterial->GetDiffuse();
	if (hitMaterial->GetAlbedoTexture() != -1)
	{
		vec4 texel = scene.GetTexture(hitMaterial->GetAlbedoTexture())->SampleBilinear(closestHit.hitTexcoord);
		materialDiffuse.r = texel.r;
		materialDiffuse.g = texel.g;
		materialDiffuse.b = texel.b;
	}

	// iOrange - if we have bump texture - then apply transformation on our normal
	if (hitMaterial->GetBumpTexture() != -1)
	{
		vec4 texel = scene.GetTexture(hitMaterial->GetBumpTexture())->SampleBilinear(closestHit.hitTexcoord);
		vec3 n = rtMath::normalize(vec3(texel.x, texel.y, texel.z) * 2.0f - 1.0f);
		closestHit.hitNormal = rtMath::transform_by_normal(n, closestHit.hitNormal);
	}

	vec3 materialSpecular = hitMaterial->GetSpecular();
	float specLevel = hitMaterial->GetSpecularLevel();
	float glossiness = hitMaterial->GetGlossiness();
	float reflectivity = hitMaterial->GetReflectivity();
	float transparency = hitMaterial->GetTransparency();

	for (int i = 0; i < numLights; ++i)
	{
		const SceneLight* light = scene.GetLight(i);
		vec3 dirToLight = light->GetPosition() - closestHit.hitPos;
		float distToLight = dirToLight.length();
		dirToLight.normalize();

		float NdotL = rtMath::dot(closestHit.hitNormal, dirToLight);
		if (NdotL <= 0.0f)
			continue;

		if (distToLight <= 0.0f)
			continue;

		Ray lightRay(closestHit.hitPos, dirToLight);

		bool inShadow = false;
		for (int j = 0; j < numShapes; ++j)
		{
			RayHitInfo shadowHit;
			const SceneShape* shape = scene.GetShape(j);
			if (shape->RayCollisionCheck(lightRay, &shadowHit, distToLight))
			{
				inShadow = true;
				break;
			}
		}

		if (!inShadow)
		{
			const vec3& lightIntensity = light->GetIntensity() * lightRay.lightStrength;
			float colorInfluence = rtMath::clamp(1.0f - reflectivity - transparency, 0.0f, 1.0f);

			// lambert
			float lambert = rtMath::clamp(NdotL, 0.0f, 1.0f);
			color.r += lambert * lightIntensity.r * materialDiffuse.r * colorInfluence;
			color.g += lambert * lightIntensity.g * materialDiffuse.g * colorInfluence;
			color.b += lambert * lightIntensity.b * materialDiffuse.b * colorInfluence;

			// point light source: sample once for specular highlight
			if (specLevel > 0.0f)
			{
				vec3 V = ray.dir;
				vec3 R = rtMath::reflect(dirToLight, closestHit.hitNormal);
				float VdotR = rtMath::dot(V, R);
				if (VdotR > 0)
				{
					float spec = powf(VdotR, glossiness) * lambert * specLevel;
					// add specular component to ray color
					color.r += spec * lightIntensity.r * materialSpecular.r * colorInfluence;
					color.g += spec * lightIntensity.g * materialSpecular.g * colorInfluence;
					color.b += spec * lightIntensity.b * materialSpecular.b * colorInfluence;
				}
			}
		}
	}

	bool bInside;
	if (rtMath::dot(closestHit.hitNormal, ray.dir) > 0.0f)
	{
		closestHit.hitNormal = -closestHit.hitNormal;
		bInside = true;
	}
	else
		bInside = false;

	if (currentDepth < m_MaximumTraceDepth)
	{
		if (reflectivity > rtMath::EPSILON && ray.lightStrength > rtMath::EPSILON)
		{
			vec3 reflDir = rtMath::reflect(ray.dir, closestHit.hitNormal);
			Ray reflRay(closestHit.hitPos + reflDir * rtMath::EPSILON, reflDir);
			//reflRay.lightStrength = ray.lightStrength * reflectivity;
			color += this->TraceRay(reflRay, scene, currentDepth + 1) * reflectivity;// * ray.lightStrength;
		}
		if (transparency > rtMath::EPSILON && ray.lightStrength > rtMath::EPSILON)
		{
			float refrIndex = bInside ? hitMaterial->GetRefractionIndex() : (1.0f / hitMaterial->GetRefractionIndex());
			float NdotI = -rtMath::dot(closestHit.hitNormal, ray.dir);
			float cosT2 = 1.0f + refrIndex * refrIndex * (NdotI * NdotI - 1.0f);
			if (cosT2 > 0.0f)
			{
				vec3 refrDir = ray.dir * refrIndex + closestHit.hitNormal * (refrIndex * NdotI - sqrtf(cosT2));
				Ray refrRay(closestHit.hitPos + refrDir * rtMath::EPSILON, refrDir);
				//refrRay.lightStrength = ray.lightStrength * transparency;
				color += this->TraceRay(refrRay, scene, currentDepth + 1) * transparency;// * ray.lightStrength;
			}
		}
	}

	color.saturate();
	return color;
}
